use alloc::{
    format,
    string::{String, ToString},
    vec,
    vec::Vec,
};
use core::{fmt, time::Duration};
use std::{
    fs::File,
    io::{self, Read, Write},
    net::{SocketAddr, UdpSocket},
};

const DEFAULT_PORT: u16 = 69;
const DEFAULT_TIMEOUT: Duration = Duration::from_secs(5);
const BLOCK_SIZE: usize = 512;
const HEADER_SIZE: usize = 4;
const MAX_PACKET: usize = HEADER_SIZE + BLOCK_SIZE;

/// Usage text shared by TFTP command front-ends.
pub const USAGE_LINES: &[&str] = &[
    "Usage: tftp <command> [args]",
    "  tftp get <server[:port]> <remote> [local]",
    "  tftp put <server[:port]> <local> [remote]",
    "  tftp help",
];

/// Parsed TFTP subcommands.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Command<'a> {
    Get {
        server: &'a str,
        remote: &'a str,
        local: Option<&'a str>,
    },
    Put {
        server: &'a str,
        local: &'a str,
        remote: Option<&'a str>,
    },
    Help,
}

/// Result of executing a TFTP command.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Response {
    Get {
        server: String,
        remote: String,
        local: String,
        bytes: usize,
        blocks: u16,
    },
    Put {
        server: String,
        local: String,
        remote: String,
        bytes: usize,
        blocks: u16,
    },
    Help,
}

/// TFTP execution error.
#[derive(Debug)]
pub enum Error {
    Io(io::Error),
    Protocol(&'static str),
    ProtocolMessage(String),
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Error::Io(err) => write!(f, "io error: {err}"),
            Error::Protocol(msg) => write!(f, "protocol error: {msg}"),
            Error::ProtocolMessage(msg) => write!(f, "protocol error: {msg}"),
        }
    }
}

impl From<io::Error> for Error {
    fn from(err: io::Error) -> Self {
        Error::Io(err)
    }
}

/// Possible parsing failures for TFTP commands.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ParseError<'a> {
    MissingSubcommand,
    MissingArgument(&'static str),
    UnknownSubcommand(&'a str),
}

impl fmt::Display for ParseError<'_> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            ParseError::MissingSubcommand => write!(f, "missing subcommand"),
            ParseError::MissingArgument(arg) => write!(f, "missing {arg} argument"),
            ParseError::UnknownSubcommand(cmd) => write!(f, "unknown subcommand '{cmd}'"),
        }
    }
}

/// Optional runtime configuration for the TFTP client.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Options {
    pub timeout: Duration,
    pub block_size: usize,
}

impl Default for Options {
    fn default() -> Self {
        Self {
            timeout: DEFAULT_TIMEOUT,
            block_size: BLOCK_SIZE,
        }
    }
}

/// Parse a TFTP command string.
pub fn parse_command<'a, I>(subcmd: &'a str, mut args: I) -> Result<Command<'a>, ParseError<'a>>
where
    I: Iterator<Item = &'a str>,
{
    match subcmd {
        "get" => {
            let server = args.next().ok_or(ParseError::MissingArgument("server"))?;
            let remote = args.next().ok_or(ParseError::MissingArgument("remote"))?;
            let local = args.next();
            Ok(Command::Get {
                server,
                remote,
                local,
            })
        }
        "put" => {
            let server = args.next().ok_or(ParseError::MissingArgument("server"))?;
            let local = args.next().ok_or(ParseError::MissingArgument("local"))?;
            let remote = args.next();
            Ok(Command::Put {
                server,
                local,
                remote,
            })
        }
        "help" => Ok(Command::Help),
        other => Err(ParseError::UnknownSubcommand(other)),
    }
}

/// Execute a parsed TFTP command.
pub fn execute(command: Command<'_>, options: Options) -> Result<Response, Error> {
    match command {
        Command::Help => Ok(Response::Help),
        Command::Get {
            server,
            remote,
            local,
        } => {
            let local_path = local.unwrap_or(remote);
            let mut client = TftpClient::new(server, options)?;
            let (bytes, blocks) = client.get(remote, local_path)?;
            Ok(Response::Get {
                server: client.server_display(),
                remote: remote.to_string(),
                local: local_path.to_string(),
                bytes,
                blocks,
            })
        }
        Command::Put {
            server,
            local,
            remote,
        } => {
            let remote_name = remote.unwrap_or_else(|| basename(local));
            let mut client = TftpClient::new(server, options)?;
            let (bytes, blocks) = client.put(local, remote_name)?;
            Ok(Response::Put {
                server: client.server_display(),
                local: local.to_string(),
                remote: remote_name.to_string(),
                bytes,
                blocks,
            })
        }
    }
}

struct TftpClient {
    socket: UdpSocket,
    server_host: String,
    server_port: u16,
    block_size: usize,
}

impl TftpClient {
    fn new(server: &str, options: Options) -> Result<Self, Error> {
        let (host, port) = parse_host(server);
        let socket = UdpSocket::bind("0.0.0.0:0").or_else(|_| UdpSocket::bind("[::]:0"))?;
        let _ = options.timeout; // axstd::net::UdpSocket currently lacks timeout configuration APIs.

        let block_size = options.block_size.clamp(8, BLOCK_SIZE);

        Ok(Self {
            socket,
            server_host: host,
            server_port: port.unwrap_or(DEFAULT_PORT),
            block_size,
        })
    }

    fn server_display(&self) -> String {
        format!("{}:{}", self.server_host, self.server_port)
    }

    fn create_request(&self, opcode: u16, filename: &str, mode: &str) -> Vec<u8> {
        let mut packet = Vec::with_capacity(HEADER_SIZE + filename.len() + mode.len() + 2);
        packet.extend_from_slice(&opcode.to_be_bytes());
        packet.extend_from_slice(filename.as_bytes());
        packet.push(0);
        packet.extend_from_slice(mode.as_bytes());
        packet.push(0);
        packet
    }

    fn create_ack(&self, block: u16) -> [u8; HEADER_SIZE] {
        let mut packet = [0u8; HEADER_SIZE];
        packet[..2].copy_from_slice(&4u16.to_be_bytes()); // Opcode::ACK
        packet[2..].copy_from_slice(&block.to_be_bytes());
        packet
    }

    fn create_data(&self, block: u16, data: &[u8]) -> Vec<u8> {
        let mut packet = Vec::with_capacity(HEADER_SIZE + data.len());
        packet.extend_from_slice(&3u16.to_be_bytes()); // Opcode::DATA
        packet.extend_from_slice(&block.to_be_bytes());
        packet.extend_from_slice(data);
        packet
    }

    fn get(&mut self, remote_file: &str, local_path: &str) -> Result<(usize, u16), Error> {
        let mut file = File::create(local_path)?;
        let request = self.create_request(1, remote_file, "octet");
        self.socket.send_to(
            &request,
            format!("{}:{}", self.server_host, self.server_port),
        )?;

        let mut expected_block: u16 = 1;
        let mut last_block: u16 = 0;
        let mut total_bytes = 0usize;
        let mut blocks = 0u16;
        let mut buffer = [0u8; MAX_PACKET];
        let mut server_tid: Option<SocketAddr> = None;

        loop {
            let (size, peer) = self.socket.recv_from(&mut buffer)?;

            if size < HEADER_SIZE {
                return Err(Error::Protocol("short packet"));
            }

            if let Some(tid) = server_tid {
                if tid != peer {
                    self.send_unknown_tid(peer)?;
                    continue;
                }
            } else {
                server_tid = Some(peer);
            }

            match u16::from_be_bytes([buffer[0], buffer[1]]) {
                3 => {
                    let block = u16::from_be_bytes([buffer[2], buffer[3]]);
                    let data = &buffer[HEADER_SIZE..size];

                    if block == expected_block {
                        file.write_all(data)?;
                        total_bytes += data.len();
                        blocks = blocks.wrapping_add(1);
                        last_block = block;
                        expected_block = expected_block.wrapping_add(1);
                    } else if block == last_block {
                        // Duplicate DATA packet; re-ACK
                    } else {
                        return Err(Error::Protocol("unexpected block number"));
                    }

                    let ack = self.create_ack(block);
                    self.socket.send_to(&ack, peer)?;

                    if data.len() < self.block_size {
                        break;
                    }
                }
                5 => {
                    let code = u16::from_be_bytes([buffer[2], buffer[3]]);
                    let message = extract_error(&buffer[HEADER_SIZE..size]);
                    return Err(Error::ProtocolMessage(format!(
                        "server error {}: {}",
                        code, message
                    )));
                }
                other => {
                    return Err(Error::ProtocolMessage(format!(
                        "unexpected opcode {}",
                        other
                    )));
                }
            }
        }

        Ok((total_bytes, blocks))
    }

    fn put(&mut self, local_path: &str, remote: &str) -> Result<(usize, u16), Error> {
        let mut file = File::open(local_path)?;
        let request = self.create_request(2, remote, "octet");
        self.socket.send_to(
            &request,
            format!("{}:{}", self.server_host, self.server_port),
        )?;

        let mut buffer = [0u8; MAX_PACKET];
        let (size, peer) = self.socket.recv_from(&mut buffer)?;
        if size < HEADER_SIZE {
            return Err(Error::Protocol("short packet"));
        }

        match u16::from_be_bytes([buffer[0], buffer[1]]) {
            4 => {
                let block = u16::from_be_bytes([buffer[2], buffer[3]]);
                if block != 0 {
                    return Err(Error::Protocol("expected ACK for block 0"));
                }
            }
            5 => {
                let code = u16::from_be_bytes([buffer[2], buffer[3]]);
                let message = extract_error(&buffer[HEADER_SIZE..size]);
                return Err(Error::ProtocolMessage(format!(
                    "server error {}: {}",
                    code, message
                )));
            }
            other => {
                return Err(Error::ProtocolMessage(format!(
                    "unexpected opcode {}",
                    other
                )));
            }
        }

        let mut block = 1u16;
        let mut bytes = 0usize;
        let mut blocks = 0u16;
        let mut data_buf = vec![0u8; self.block_size];

        loop {
            let read = file.read(&mut data_buf)?;
            let data = &data_buf[..read];

            let packet = self.create_data(block, data);
            self.socket.send_to(&packet, peer)?;

            loop {
                let (size, ack_peer) = self.socket.recv_from(&mut buffer)?;
                if ack_peer != peer {
                    self.send_unknown_tid(ack_peer)?;
                    continue;
                }

                if size < HEADER_SIZE {
                    return Err(Error::Protocol("short packet"));
                }

                match u16::from_be_bytes([buffer[0], buffer[1]]) {
                    4 => {
                        let ack_block = u16::from_be_bytes([buffer[2], buffer[3]]);
                        if ack_block == block {
                            break;
                        } else if ack_block == block.wrapping_sub(1) {
                            self.socket.send_to(&packet, peer)?;
                            continue;
                        } else {
                            return Err(Error::Protocol("unexpected ACK block"));
                        }
                    }
                    5 => {
                        let code = u16::from_be_bytes([buffer[2], buffer[3]]);
                        let message = extract_error(&buffer[HEADER_SIZE..size]);
                        return Err(Error::ProtocolMessage(format!(
                            "server error {}: {}",
                            code, message
                        )));
                    }
                    other => {
                        return Err(Error::ProtocolMessage(format!(
                            "unexpected opcode {}",
                            other
                        )));
                    }
                }
            }

            bytes += data.len();
            blocks = blocks.wrapping_add(1);

            if data.len() < self.block_size {
                break;
            }
            block = block.wrapping_add(1);
        }

        Ok((bytes, blocks))
    }

    fn send_unknown_tid(&self, addr: SocketAddr) -> Result<(), Error> {
        let mut packet = Vec::with_capacity(HEADER_SIZE + 16);
        packet.extend_from_slice(&0x0005u16.to_be_bytes());
        packet.extend_from_slice(&0x0005u16.to_be_bytes());
        packet.extend_from_slice(b"Unknown transfer ID");
        packet.push(0);
        let _ = self.socket.send_to(&packet, addr);
        Ok(())
    }
}

fn parse_host(server: &str) -> (String, Option<u16>) {
    if let Some(end) = server.find(']') {
        if server.starts_with('[') {
            let host = server[1..end].to_string();
            let port = server[end + 1..]
                .strip_prefix(':')
                .and_then(|s| (!s.is_empty()).then_some(s))
                .and_then(|s| s.parse().ok());
            return (host, port);
        }
    }

    if let Some(pos) = server.rfind(':') {
        let head = &server[..pos];
        let tail = &server[pos + 1..];
        if head.contains(':') {
            return (server.to_string(), None);
        }
        if !tail.is_empty() && tail.chars().all(|c| c.is_ascii_digit()) {
            if let Ok(port) = tail.parse() {
                return (head.to_string(), Some(port));
            }
        }
    }

    (server.to_string(), None)
}

fn basename(path: &str) -> &str {
    path.rsplit_once('/').map(|(_, name)| name).unwrap_or(path)
}

fn extract_error(bytes: &[u8]) -> String {
    let len = bytes.iter().position(|&b| b == 0).unwrap_or(bytes.len());
    String::from_utf8_lossy(&bytes[..len]).into_owned()
}
